/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

package org.elasticsearch.xpack.esql.expression.function.scalar.conditional;

import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.compute.data.Block;
import org.elasticsearch.compute.data.BooleanBlock;
import org.elasticsearch.compute.data.ElementType;
import org.elasticsearch.compute.data.Page;
import org.elasticsearch.compute.operator.DriverContext;
import org.elasticsearch.compute.operator.EvalOperator;
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
import org.elasticsearch.core.Releasable;
import org.elasticsearch.core.Releasables;
import org.elasticsearch.xpack.esql.core.expression.Expression;
import org.elasticsearch.xpack.esql.core.expression.Literal;
import org.elasticsearch.xpack.esql.core.expression.Nullability;
import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
import org.elasticsearch.xpack.esql.core.tree.Source;
import org.elasticsearch.xpack.esql.core.type.DataType;
import org.elasticsearch.xpack.esql.expression.function.Example;
import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
import org.elasticsearch.xpack.esql.expression.function.Param;
import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction;
import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
import org.elasticsearch.xpack.esql.planner.PlannerUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
import static org.elasticsearch.xpack.esql.core.type.DataType.NULL;

public final class Case extends EsqlScalarFunction {
    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Case", Case::new);

    record Condition(Expression condition, Expression value) {}

    private final List<Condition> conditions;
    private final Expression elseValue;
    private DataType dataType;

    @FunctionInfo(
        returnType = {
            "boolean",
            "cartesian_point",
            "date",
            "double",
            "geo_point",
            "integer",
            "ip",
            "keyword",
            "long",
            "text",
            "unsigned_long",
            "version" },
        description = """
            Accepts pairs of conditions and values. The function returns the value that
            belongs to the first condition that evaluates to `true`.

            If the number of arguments is odd, the last argument is the default value which
            is returned when no condition matches. If the number of arguments is even, and
            no condition matches, the function returns `null`.""",
        examples = {
            @Example(description = "Determine whether employees are monolingual, bilingual, or polyglot:", file = "docs", tag = "case"),
            @Example(
                description = "Calculate the total connection success rate based on log messages:",
                file = "conditional",
                tag = "docsCaseSuccessRate"
            ),
            @Example(
                description = "Calculate an hourly error rate as a percentage of the total number of log messages:",
                file = "conditional",
                tag = "docsCaseHourlyErrorRate"
            ) }
    )
    public Case(
        Source source,
        @Param(name = "condition", type = { "boolean" }, description = "A condition.") Expression first,
        @Param(
            name = "trueValue",
            type = {
                "boolean",
                "cartesian_point",
                "date",
                "double",
                "geo_point",
                "integer",
                "ip",
                "keyword",
                "long",
                "text",
                "unsigned_long",
                "version" },
            description = "The value that's returned when the corresponding condition is the first to evaluate to `true`. "
                + "The default value is returned when no condition matches."
        ) List<Expression> rest
    ) {
        super(source, Stream.concat(Stream.of(first), rest.stream()).toList());
        int conditionCount = children().size() / 2;
        conditions = new ArrayList<>(conditionCount);
        for (int c = 0; c < conditionCount; c++) {
            conditions.add(new Condition(children().get(c * 2), children().get(c * 2 + 1)));
        }
        elseValue = elseValueIsExplicit() ? children().get(children().size() - 1) : new Literal(source, null, NULL);
    }

    private Case(StreamInput in) throws IOException {
        this(
            Source.readFrom((PlanStreamInput) in),
            in.readNamedWriteable(Expression.class),
            in.readNamedWriteableCollectionAsList(Expression.class)
        );
    }

    @Override
    public void writeTo(StreamOutput out) throws IOException {
        source().writeTo(out);
        out.writeNamedWriteable(children().get(0));
        out.writeNamedWriteableCollection(children().subList(1, children().size()));
    }

    @Override
    public String getWriteableName() {
        return ENTRY.name;
    }

    private boolean elseValueIsExplicit() {
        return children().size() % 2 == 1;
    }

    @Override
    public DataType dataType() {
        if (dataType == null) {
            resolveType();
        }
        return dataType;
    }

    @Override
    protected TypeResolution resolveType() {
        if (childrenResolved() == false) {
            return new TypeResolution("Unresolved children");
        }

        if (children().size() < 2) {
            return new TypeResolution(format(null, "expected at least two arguments in [{}] but got {}", sourceText(), children().size()));
        }

        for (int c = 0; c < conditions.size(); c++) {
            Condition condition = conditions.get(c);

            TypeResolution resolution = TypeResolutions.isBoolean(
                condition.condition,
                sourceText(),
                TypeResolutions.ParamOrdinal.fromIndex(c * 2)
            );
            if (resolution.unresolved()) {
                return resolution;
            }

            resolution = resolveValueType(condition.value, c * 2 + 1);
            if (resolution.unresolved()) {
                return resolution;
            }
        }

        return resolveValueType(elseValue, conditions.size() * 2);
    }

    private TypeResolution resolveValueType(Expression value, int position) {
        if (dataType == null || dataType == NULL) {
            dataType = value.dataType();
            return TypeResolution.TYPE_RESOLVED;
        }
        return TypeResolutions.isType(
            value,
            t -> t == dataType,
            sourceText(),
            TypeResolutions.ParamOrdinal.fromIndex(position),
            dataType.typeName()
        );
    }

    @Override
    public Nullability nullable() {
        return Nullability.UNKNOWN;
    }

    @Override
    public Expression replaceChildren(List<Expression> newChildren) {
        return new Case(source(), newChildren.get(0), newChildren.subList(1, newChildren.size()));
    }

    @Override
    protected NodeInfo<? extends Expression> info() {
        return NodeInfo.create(this, Case::new, children().get(0), children().subList(1, children().size()));
    }

    @Override
    public boolean foldable() {
        for (Condition condition : conditions) {
            if (condition.condition.foldable() == false) {
                return false;
            }
            Boolean b = (Boolean) condition.condition.fold();
            if (b != null && b) {
                return condition.value.foldable();
            }
        }
        return elseValue.foldable();
    }

    @Override
    public Object fold() {
        for (Condition condition : conditions) {
            Boolean b = (Boolean) condition.condition.fold();
            if (b != null && b) {
                return condition.value.fold();
            }
        }
        return elseValue.fold();
    }

    /**
     * Fold the arms of {@code CASE} statements.
     * <ol>
     *     <li>
     *         Conditions that evaluate to {@code false} are removed so
     *         {@code EVAL c=CASE(false, foo, b, bar, bort)} becomes
     *         {@code EVAL c=CASE(b, bar, bort)}.
     *     </li>
     *     <li>
     *         Conditions that evaluate to {@code true} stop evaluation and
     *         return themselves so {@code EVAL c=CASE(true, foo, bar)} becomes
     *         {@code EVAL c=foo}.
     *     </li>
     * </ol>
     * And those two combine so {@code EVAL c=CASE(false, foo, b, bar, true, bort, el)} becomes
     * {@code EVAL c=CASE(b, bar, bort)}.
     */
    public Expression partiallyFold() {
        List<Expression> newChildren = new ArrayList<>(children().size());
        boolean modified = false;
        for (Condition condition : conditions) {
            if (condition.condition.foldable() == false) {
                newChildren.add(condition.condition);
                newChildren.add(condition.value);
                continue;
            }
            modified = true;
            Boolean b = (Boolean) condition.condition.fold();
            if (b != null && b) {
                newChildren.add(condition.value);
                return finishPartialFold(newChildren);
            }
        }
        if (modified == false) {
            return this;
        }
        if (elseValueIsExplicit()) {
            newChildren.add(elseValue);
        }
        return finishPartialFold(newChildren);
    }

    private Expression finishPartialFold(List<Expression> newChildren) {
        if (newChildren.size() == 1) {
            return newChildren.get(0);
        }
        return replaceChildren(newChildren);
    }

    @Override
    public ExpressionEvaluator.Factory toEvaluator(Function<Expression, ExpressionEvaluator.Factory> toEvaluator) {
        ElementType resultType = PlannerUtils.toElementType(dataType());
        List<ConditionEvaluatorSupplier> conditionsFactories = conditions.stream()
            .map(c -> new ConditionEvaluatorSupplier(toEvaluator.apply(c.condition), toEvaluator.apply(c.value)))
            .toList();
        ExpressionEvaluator.Factory elseValueFactory = toEvaluator.apply(elseValue);
        return new ExpressionEvaluator.Factory() {
            @Override
            public ExpressionEvaluator get(DriverContext context) {
                return new CaseEvaluator(
                    context,
                    resultType,
                    conditionsFactories.stream().map(x -> x.apply(context)).toList(),
                    elseValueFactory.get(context)
                );
            }

            @Override
            public String toString() {
                return "CaseEvaluator[resultType="
                    + resultType
                    + ", conditions="
                    + conditionsFactories
                    + ", elseVal="
                    + elseValueFactory
                    + ']';
            }
        };
    }

    record ConditionEvaluatorSupplier(ExpressionEvaluator.Factory condition, ExpressionEvaluator.Factory value)
        implements
            Function<DriverContext, ConditionEvaluator> {
        @Override
        public ConditionEvaluator apply(DriverContext driverContext) {
            return new ConditionEvaluator(condition.get(driverContext), value.get(driverContext));
        }

        @Override
        public String toString() {
            return "ConditionEvaluator[" + "condition=" + condition + ", value=" + value + ']';
        }
    }

    record ConditionEvaluator(EvalOperator.ExpressionEvaluator condition, EvalOperator.ExpressionEvaluator value) implements Releasable {
        @Override
        public void close() {
            Releasables.closeExpectNoException(condition, value);
        }
    }

    private record CaseEvaluator(
        DriverContext driverContext,
        ElementType resultType,
        List<ConditionEvaluator> conditions,
        EvalOperator.ExpressionEvaluator elseVal
    ) implements EvalOperator.ExpressionEvaluator {
        @Override
        public Block eval(Page page) {
            /*
             * We have to evaluate lazily so any errors or warnings that would be
             * produced by the right hand side are avoided. And so if anything
             * on the right hand side is slow we skip it.
             *
             * And it'd be good if that lazy evaluation were fast. But this
             * implementation isn't. It's fairly simple - running position at
             * a time - but it's not at all fast.
             */
            int positionCount = page.getPositionCount();
            try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) {
                position: for (int p = 0; p < positionCount; p++) {
                    int[] positions = new int[] { p };
                    Page limited = new Page(
                        IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new)
                    );
                    try (Releasable ignored = limited::releaseBlocks) {
                        for (ConditionEvaluator condition : conditions) {
                            try (BooleanBlock b = (BooleanBlock) condition.condition.eval(limited)) {
                                if (b.isNull(0)) {
                                    continue;
                                }
                                if (false == b.getBoolean(b.getFirstValueIndex(0))) {
                                    continue;
                                }
                                try (Block values = condition.value.eval(limited)) {
                                    result.copyFrom(values, 0, 1);
                                    continue position;
                                }
                            }
                        }
                        try (Block values = elseVal.eval(limited)) {
                            result.copyFrom(values, 0, 1);
                        }
                    }
                }
                return result.build();
            }
        }

        @Override
        public void close() {
            Releasables.closeExpectNoException(() -> Releasables.close(conditions), elseVal);
        }

        @Override
        public String toString() {
            return "CaseEvaluator[resultType=" + resultType + ", conditions=" + conditions + ", elseVal=" + elseVal + ']';
        }
    }
}
